import keras
import tensorflow as tf
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from model import Model
from preprocess import preprocess
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # 这一行注释掉就是使用gpu，不注释就是使用cpu


if __name__ == 'main':
    temp, x, y, x_test, y_test = preprocess()  # temp is the total
    model = Model()
    # Compile the model
    model.compile(optimizer='adam', loss='mean_squared_error')

    # Train the model
    model.fit(x, y, batch_size=1, epochs=1)

    predictions = model.predict(x_test)
    predictions = scaler.inverse_transform(predictions)

    # Get the root mean squared error (RMSE)
    rmse = np.sqrt(np.mean(((predictions - y_test) ** 2)))
    print(rmse)

    train = temp[:training_data_len]
    valid = temp[training_data_len:]
    # predictions
    # Visualize the data
    plt.figure(figsize=(16, 6))
    plt.title('Model')
    plt.xlabel('Date', fontsize=18)
    plt.ylabel('Recovered', fontsize=18)
    plt.plot(temp['ObservationDate'][:374], temp['Recovered'][:374])
    plt.plot(temp['ObservationDate'][374:], valid)
    plt.plot(temp['ObservationDate'][374:], predictions)
    plt.legend(['Train', 'Val', 'Predictions'], loc='lower right')
    plt.savefig('prediction.jpg')
